﻿using System;
using System.Collections.Generic;
using System.Disposables;
using System.Linq;

namespace RxPowerToys.ObservableExtensions
{
    public static class ConcurrencyLimitingMerge
    {
        /// <summary>
        /// Acts like Observable.Merge() but limits the number of concurrently active inner observables.
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="asyncInput">Enumerable sequence of source observables</param>
        /// <param name="maxConcurrent">Maxmum number of concurrently active observables.</param>
        /// <returns></returns>
        public static IObservable<T> Merge<T>(this IEnumerable<IObservable<T>> asyncInput, int maxConcurrent)
        {
            return Observable.CreateWithDisposable<T>(o =>
            {
                object lockToken = new Object();
                var cancelBool = new BooleanDisposable();
                var cancelGroup = new CompositeDisposable(cancelBool);
                var inputEnum = asyncInput.GetEnumerator();
                Action startNext = null;
                bool inputAtEnd = false;
                Exception failure = null;
                startNext = () =>
                {
                    lock (lockToken)
                    {
                        while (cancelGroup.Count <= maxConcurrent)
                        {
                            if (cancelBool.IsDisposed || failure != null || !inputEnum.MoveNext())
                            {
                                inputAtEnd = true;
                                break;
                            }
                            var subscription = new FutureDisposable();
                            cancelGroup.Add(subscription);
                            subscription.Set(
                                inputEnum.Current.Subscribe(o.OnNext,
                                    e =>
                                    {
                                        lock (lockToken)
                                        {
                                            cancelGroup.Remove(subscription);
                                            if (failure == null)
                                                failure = e;
                                        }
                                        startNext();
                                    },
                                    () =>
                                    {
                                        lock (lockToken)
                                        {
                                            cancelGroup.Remove(subscription);
                                        }
                                        startNext();
                                    }
                                )
                            );
                        }
                        if (inputAtEnd && cancelGroup.Count == 1 && !cancelBool.IsDisposed)
                        {
                            if (failure == null)
                                o.OnCompleted();
                            else
                                o.OnError(failure);
                        }
                    }
                };
                startNext();
                return cancelGroup;
            });
        }
    }
}
